from string import Template

class PromptBuilder:
    """构建 GPT 的 prompt（包含单条和批量两种）"""

    def __init__(self):
        pass

    def build(self, table_name: str, column_samples: dict) -> str:
        formatted = "\n".join(f"- {col}: {samples[:5]}" for col, samples in column_samples.items())
        return f"""You are a data analyst helping recover the original table and column names from a hashed database schema.
Table name: `{table_name}`
Column samples:
{formatted}

Please respond with a JSON object:
{{
  "table": "realistic_table_name",
  "columns": {{
    "hashed_col1": "recovered_column_name1",
    "hashed_col2": "recovered_column_name2"
  }}
}}
Be concise.
"""

    def build_single_fk(self, t1, c1, t2, c2, set_1=None, set_2=None) -> str:
        vals1 = list(set_1)[:20] if set_1 else "N/A"
        vals2 = list(set_2)[:20] if set_2 else "N/A"
        return f"""
You are an expert database architect specializing in banking and financial systems.
For the following two columns, decide whether the first column references the second (i.e., is a foreign key -> primary key relationship).
Provide a JSON object with fields: decision (true/false), child (table.col), parent (table.col), confidence (0.0-1.0), reason.

Table1.Column1: `{t1}.{c1}`, sample values: {vals1}
Table2.Column2: `{t2}.{c2}`, sample values: {vals2}

Example output:
{{
  "decision": true,
  "child": "orders.user_id",
  "parent": "users.id",
  "confidence": 0.9,
  "reason": "column name contains table name + value overlap"
}}
If uncertain, set decision to false.
""".strip()

    def build_batch_fk(self, batch):
        """
        batch: list of tuples (t1,c1,t2,c2)
        返回一个 prompt，要求 GPT 批量输出 JSON: {"1": {...}, "2": {...}, ...}
        每个 value 是一个对象 {decision, child, parent, confidence, reason}
        """
        lines = [
            "You are an expert database architect specializing in banking and financial systems.",
            "For each candidate below, decide whether it is a foreign key relationship and indicate the direction.",
            "Return a JSON object where keys are the candidate indices (1-based) and values are objects with fields: decision (true/false), child (table.col), parent (table.col), confidence (0.0-1.0), reason."
        ]
        for idx, (t1, c1, t2, c2, _) in enumerate(batch, start=1):
            lines.append(f"{idx}. {t1}.{c1}  <->  {t2}.{c2}")
        lines.append('\nReturn JSON like: { "1": { "decision": true, "child":"orders.user_id", "parent":"users.id", "confidence":0.9, "reason":"..." }, "2": { ... } }')
        return "\n".join(lines)